package hydrap2p.helper;

import java.nio.ByteBuffer;
import java.util.Timer;
import java.util.TimerTask;

import hydrap2p.Service;
import hydrap2p.ServiceLocation;

import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TField;
import org.apache.thrift.protocol.TList;
import org.apache.thrift.protocol.TMap;
import org.apache.thrift.protocol.TMessage;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TSet;
import org.apache.thrift.protocol.TStruct;
import org.apache.thrift.transport.TTransportException;

/**
 * Implements a fake Protocol that wraps around a Binary Protocol
 * 
 * It uses PooledProtocolManager to take exclusive control of a connection and later return it
 * 
 * Thrift Client classes use protocols in one of two ways:
 * 	protocol.writeMessageBegin(); protocol.writeMessageEnd(); protocol.readMessageBegin(); protocol.readMessageEnd();
 * OR
 * 	protocol.writeMessageBegin(); protocol.writeMessageEnd();
 * 
 * In the first case, we can return the connection during readMessageEnd();
 * To handle the second case, we start a timer to wait for a certain time before releasing the connection.
 * 
 * @author Hufman
 *
 */
public class TPooledProtocol extends TProtocol {

	private ServiceLocation location;
	private TProtocol protocol;
	private Timer detector;
	private static final int DETECTTIMEOUT=1000;	// number of milliseconds to wait after writeMessageEnd() while not seeing readMessageBegin() to release the connection

	protected TPooledProtocol(ServiceLocation location) throws TTransportException {
		super(null);
		this.location=location;
		try {
			trans_ = PooledProtocolManager.checkout(location);
			protocol=new TBinaryProtocol(trans_);
			PooledProtocolManager.checkin(location, trans_);
			trans_ = null;
			protocol = null;
		}
		catch (TTransportException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}
	

	@Override
	public void writeMessageBegin(TMessage message) throws TException {
		writeMessageBegin(message, 5);
	}

	public void writeMessageBegin(TMessage message, int tries) throws TException {
		try {
			trans_ = PooledProtocolManager.checkout(location);
			protocol=new TBinaryProtocol(trans_);
			protocol.writeMessageBegin(message);
			trans_.flush();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			if (tries>0)
			{
				try {
					writeMessageBegin(message, tries-1);
				}
				catch (TException e1)
				{
					throw e1;
				}
			}
			else
				throw e;
		}
	}

	@Override
	public void writeMessageEnd() throws TException {
		try {
			protocol.writeMessageEnd();
			detector=new Timer();
			detector.schedule(new DetectOneWay(), DETECTTIMEOUT);
		}
		catch (TException e)
		{
			detector.cancel();
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public TMessage readMessageBegin() throws TException {
		try {
			detector.cancel();
			return protocol.readMessageBegin();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void readMessageEnd() throws TException {
		try {
			protocol.readMessageEnd();
			PooledProtocolManager.checkin(location, trans_);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}
	
	private class DetectOneWay extends TimerTask
	{
		@Override
		public void run() {
			PooledProtocolManager.close(location, trans_);
			detector.cancel();
		}
	}

	@Override
	public ByteBuffer readBinary() throws TException {
		try {
			return protocol.readBinary();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public boolean readBool() throws TException {
		try {
			return protocol.readBool();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public byte readByte() throws TException {
		try {
			return protocol.readByte();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public double readDouble() throws TException {
		try {
			return protocol.readDouble();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public TField readFieldBegin() throws TException {
		try {	
			return protocol.readFieldBegin();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void readFieldEnd() throws TException {
		try {
			protocol.readFieldEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public short readI16() throws TException {
		try {
			return protocol.readI16();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public int readI32() throws TException {
		try {
			return protocol.readI32();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public long readI64() throws TException {
		try {
			return protocol.readI64();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public TList readListBegin() throws TException {
		try {
			return protocol.readListBegin();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void readListEnd() throws TException {
		try {
			protocol.readListEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public TMap readMapBegin() throws TException {
		try {
			return protocol.readMapBegin();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void readMapEnd() throws TException {
		try {
			protocol.readMapEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public TSet readSetBegin() throws TException {
		try {
			return protocol.readSetBegin();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void readSetEnd() throws TException {
		try {
			protocol.readSetEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public String readString() throws TException {
		try {
			return protocol.readString();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public TStruct readStructBegin() throws TException {
		try {
			return protocol.readStructBegin();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void readStructEnd() throws TException {
		try {
			protocol.readStructEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeBinary(ByteBuffer bin) throws TException {
		try {
			protocol.writeBinary(bin);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeBool(boolean b) throws TException {
		try {
			protocol.writeBool(b);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeByte(byte b) throws TException {
		try {
			protocol.writeByte(b);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeDouble(double dub) throws TException {
		try {
			protocol.writeDouble(dub);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeFieldBegin(TField field) throws TException {
		try {
			protocol.writeFieldBegin(field);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeFieldEnd() throws TException {
		try {
			protocol.writeFieldEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeFieldStop() throws TException {
		try {
			protocol.writeFieldStop();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeI16(short i16) throws TException {
		try {
			protocol.writeI16(i16);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeI32(int i32) throws TException {
		try {
			protocol.writeI32(i32);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeI64(long i64) throws TException {
		try {
			protocol.writeI64(i64);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeListBegin(TList list) throws TException {
		try {
			protocol.writeListBegin(list);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeListEnd() throws TException {
		try {
			protocol.writeListEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeMapBegin(TMap map) throws TException {
		try {
			protocol.writeMapBegin(map);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeMapEnd() throws TException {
		try {
			protocol.writeMapEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeSetBegin(TSet set) throws TException {
		try {
			protocol.writeSetBegin(set);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeSetEnd() throws TException {
		try {
			protocol.writeSetEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeString(String str) throws TException {
		try {
			protocol.writeString(str);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeStructBegin(TStruct struct) throws TException {
		try {
			protocol.writeStructBegin(struct);
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

	@Override
	public void writeStructEnd() throws TException {
		try {
			protocol.writeStructEnd();
		}
		catch (TException e)
		{
			PooledProtocolManager.close(location, trans_);
			throw e;
		}
	}

}
